import json
from time import time

import pandas as pd
import numpy as np
import stomp
from sklearn.ensemble import GradientBoostingRegressor


class MachineTaskListener(object):
    def __init__(self, host, port):
        self.conn = stomp.Connection10([(host, port)], auto_content_length=False)
        self.conn.connect()

        x = pd.DataFrame(np.random.random(2000).reshape(-1, 20))
        y = pd.Series(np.random.random(100))
        self.model = GradientBoostingRegressor().fit(x, y)

    def on_message(self, headers, msg):
        t0 = time()
        print("received message %s" % json.loads(msg))
        task_message = json.loads(msg)
        task_param = json.loads(task_message['TaskParam'])
        x_test = pd.DataFrame(np.array(task_param['X_Test']).reshape(-1, 20))
        task_message['TaskResult'] = self.model.predict(x_test)[0]
        task_message['TotalSeconds'] = time() - t0

        replay_msg = json.dumps(task_message)
        self.conn.send(headers['reply-to'], replay_msg)

    def on_error(self, headers, message):
        print('received an error %s' % message)
